# file: download_video.py
import os
import sys
import time
import datetime
import shutil
import subprocess
from ctypes import *
import glob
import concurrent.futures
import threading

# ============================================================================
# 1. 配置参数 (已更新)
# ============================================================================
NVR_IP = b"10.11.36.3"
NVR_PORT = 80
NVR_USERNAME = b"user"
NVR_PASSWORD = b"123456"
# 备用凭据:
# NVR_USERNAME = b"root"
# NVR_PASSWORD = b"Nanjing@8625"

CAMERAS = [
    {"code": "210235C8Q33254000011", "channel_id": 1},
    {"code": "210235C8Q33255000005", "channel_id": 2},
    {"code": "210235C8Q33255000006", "channel_id": 3},
]
DOWNLOAD_DURATION_SECONDS = 60
LOCAL_VIDEO_SAVE_DIR = "./downloaded_videos_temp"
MAX_PROCESSING_WORKERS = 8 
FINAL_IMAGE_ROOT = "/home/smf/image" # 注意：此路径下的抽帧功能已禁用
FINAL_TS_ROOT = "/mnt/nfs_share/"
SDK_LIB_PATH = "/root/egs_nvr-master_x86/lib"

# ============================================================================
# 2. SDK 初始化和定义 (MODIFIED: 增加错误码获取)
# ============================================================================
net_sdk_dll_path = os.path.join(SDK_LIB_PATH, "libNetDEVSDK.so")
os.environ['LD_LIBRARY_PATH'] = f"{SDK_LIB_PATH}:{os.environ.get('LD_LIBRARY_PATH', '')}"
try:
    from netdevsdk import *
except ImportError:
    print("警告：无法从'netdevsdk'导入类型定义，将使用占位符。")
    class c_void_p: pass
    class POINTER: pass
    class c_int: pass
    class c_char_p: pass
    class c_longlong: pass
    class c_int64: pass
    class c_int32: pass
    class byref: pass
    class pointer: pass
    def load_library(path): return cdll.LoadLibrary(path)
    class NETDEV_DEVICE_LOGIN_INFO_S: pass
    class NETDEV_SELOG_INFO_S: pass
    class NETDEV_PLAYBACKCOND_S: pass
    class NETDEV_VIDEO_CHL_DETAIL_INFO_S: pass

class DevClass:
    def __init__(self):
        self.egsSdk = self.load_sdk()
        if self.egsSdk: self.setup_sdk_functions()
        self.iUserID = -1

    def load_sdk(self):
        if not os.path.exists(net_sdk_dll_path):
            print(f"错误：SDK动态库文件不存在于路径: {net_sdk_dll_path}")
            print("请检查SDK_LIB_PATH配置是否正确，以及文件是否存在。")
            return None
        try:
            egs_sdk = load_library(net_sdk_dll_path)
            print(f"动态库[{net_sdk_dll_path}]加载成功")
            return egs_sdk
        except Exception as e:
            print(f"错误：动态库[{net_sdk_dll_path}]加载失败: {e}")
            print("提示：在无图形化界面的Linux服务器上，这通常是由于缺少依赖库导致的。")
            print(f"     请尝试在服务器上运行命令 'ldd {net_sdk_dll_path}' 来检查是否有 'not found' 的项目。")
            return None

    def setup_sdk_functions(self):
        self.egsSdk.NETDEV_Login_V30.restype = c_longlong
        self.egsSdk.NETDEV_GetFileByTime.argtypes = [c_longlong, POINTER(NETDEV_PLAYBACKCOND_S), c_char_p, c_int]
        self.egsSdk.NETDEV_GetFileByTime.restype = c_longlong
        self.egsSdk.NETDEV_Logout.argtypes = [c_longlong]
        self.egsSdk.NETDEV_PlayBackControl.argtypes = [c_longlong, c_int, POINTER(c_int64)]
        self.egsSdk.NETDEV_StopGetFile.argtypes = [c_longlong]
        # NEW: 增加获取最后错误码的函数定义
        self.egsSdk.NETDEV_GetLastError.restype = c_int

    # NEW: 新增一个方法用于调用SDK的错误码函数
    def get_last_error(self):
        if self.egsSdk and hasattr(self.egsSdk, 'NETDEV_GetLastError'):
            return self.egsSdk.NETDEV_GetLastError()
        return -1 # 返回-1如果SDK本身都未加载或函数不存在

    def sdk_init(self):
        if self.egsSdk.NETDEV_Init() != 0:
            # MODIFIED: 在失败时，调用新函数打印详细错误码
            error_code = self.get_last_error()
            print(f"错误：初始化SDK失败，SDK返回的错误码是: {error_code}")
            return False
        else:
            print("初始化SDK成功")
            return True

    def login(self, ip, port, username, password):
        login_info = NETDEV_DEVICE_LOGIN_INFO_S()
        login_info.szIPAddr = ip
        login_info.dwPort = port
        login_info.szUserName = username
        login_info.szPassword = password
        login_info.dwLoginProto = 0  # TCP
        login_info.dwDeviceType = 0
        se_log_info = NETDEV_SELOG_INFO_S()
        
        self.iUserID = self.egsSdk.NETDEV_Login_V30(pointer(login_info), pointer(se_log_info))
        
        if self.iUserID <= 0:
            error_code = self.get_last_error()
            print(f"错误：登录NVR失败 (IP: {ip.decode()})，SDK返回的错误码是: {error_code}")
            return False
        
        print(f"登录NVR成功 (IP: {ip.decode()})")
        return True
    
    def download(self, channel_id, begin_time, end_time, save_filename):
        os.makedirs(os.path.dirname(save_filename), exist_ok=True)
        playback_cond = NETDEV_PLAYBACKCOND_S()
        playback_cond.dwChannelID = channel_id
        playback_cond.tBeginTime = begin_time
        playback_cond.tEndTime = end_time
        playback_cond.dwLinkMode = 1
        playback_cond.dwDownloadSpeed = 8
        
        get_file = self.egsSdk.NETDEV_GetFileByTime(self.iUserID, byref(playback_cond), save_filename.encode('utf-8'), 0)
        if get_file <= 0:
            print(f"通道 {channel_id} 启动下载失败")
            return False
        
        print(f"通道 {channel_id} 开始下载...")
        download_completed = False
        timeout_start = time.time()
        max_wait_time = DOWNLOAD_DURATION_SECONDS + 30
        
        try:
            play_time = c_int64(0)
            while True:
                if time.time() - timeout_start > max_wait_time:
                    print(f"通道 {channel_id} 下载超时")
                    break
                
                if self.egsSdk.NETDEV_PlayBackControl(get_file, 3, byref(play_time)) == -1:
                    if play_time.value >= end_time:
                        download_completed = True
                    break
                
                if play_time.value >= end_time:
                    download_completed = True
                    break
                    
                time.sleep(1)
        finally:
            self.egsSdk.NETDEV_StopGetFile(get_file)
            
        return download_completed

    def logout(self):
        if self.iUserID > 0 and self.egsSdk.NETDEV_Logout(self.iUserID) != 0:
            print("登出失败")
        else:
            print("登出成功")

    def sdk_cleanup(self):
        if self.egsSdk.NETDEV_Cleanup() != 0:
            print(f"释放sdk失败")
        else:
            print("释放sdk成功")

    def __del__(self):
        self.logout()
        self.sdk_cleanup()
# ============================================================================
# 3. 视频处理函数 (已根据要求修改)
# ============================================================================

# MODIFIED: 此功能已根据您的要求禁用，不再保存逐帧图片
# def extract_images_gpu(source_mp4_path, camera_code, start_timestamp):
#     """(GPU加速)将MP4视频抽帧为JPG图片。"""
#     print(f"\n--- [GPU处理] {os.path.basename(source_mp4_path)}: 开始抽帧为JPG ---")
#     if not os.path.exists(source_mp4_path):
#         print(f"[处理失败] 源文件不存在: {source_mp4_path}")
#         return
#     try:
#         start_dt = datetime.datetime.fromtimestamp(start_timestamp)
#         final_dir = os.path.join(FINAL_IMAGE_ROOT, camera_code, start_dt.strftime('%Y/%m/%d/%H'), start_dt.strftime('%H%M%S'))
#         os.makedirs(final_dir, exist_ok=True)
#         
#         output_pattern = os.path.join(final_dir, "%06d.jpg")
#         
#         ffmpeg_command = [
#             "ffmpeg", "-y",
#             "-hwaccel", "nvdec",
#             "-i", source_mp4_path,
#             "-q:v", "2",
#             output_pattern
#         ]
#         
#         subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
#         print(f"[GPU处理成功] {os.path.basename(source_mp4_path)} -> {final_dir}")
#     except Exception as e:
#         print(f"[GPU处理错误] {os.path.basename(source_mp4_path)} 抽帧异常: {e.stderr if hasattr(e, 'stderr') else e}")


def create_jpg_thumbnail_from_video_gpu(source_mp4_path, camera_code, start_timestamp):
    """(GPU加速)直接从视频高效生成JPG缩略图。"""
    print(f"--- [GPU处理] {os.path.basename(source_mp4_path)}: 开始生成缩略图 ---")
    if not os.path.exists(source_mp4_path):
        print(f"[处理失败] 源文件不存在: {source_mp4_path}")
        return

    try:
        start_dt = datetime.datetime.fromtimestamp(start_timestamp)
        image_folder_name = start_dt.strftime('%H%M%S')
        dest_dir = os.path.join(FINAL_TS_ROOT, camera_code, start_dt.strftime('%Y/%m/%d'), 'source')
        os.makedirs(dest_dir, exist_ok=True)
        
        jpg_filename = image_folder_name + "_60.jpg"
        dest_jpg_path = os.path.join(dest_dir, jpg_filename)

        ffmpeg_command = [
            "ffmpeg", "-y",
            "-hwaccel", "nvdec",
            "-i", source_mp4_path,
            "-ss", "00:00:00.1",
            "-vframes", "1",
            "-q:v", "2",
            dest_jpg_path
        ]
        
        subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
        print(f"[GPU处理成功] 缩略图已保存到: {dest_jpg_path}")
    except Exception as e:
        print(f"[GPU处理错误] {os.path.basename(source_mp4_path)} 生成缩略图异常: {e.stderr if hasattr(e, 'stderr') else e}")


def convert_mp4_to_ts(source_mp4_path, camera_code, start_timestamp):
    """将MP4视频分割为TS文件。"""
    print(f"--- [流复制] {os.path.basename(source_mp4_path)}: 分割为TS ---")
    if not os.path.exists(source_mp4_path): return []
    
    start_dt = datetime.datetime.fromtimestamp(start_timestamp)
    final_dir = os.path.join(FINAL_TS_ROOT, camera_code, start_dt.strftime('%Y/%m/%d'), start_dt.strftime('%H'))
    final_ts_paths = []
    
    try:
        os.makedirs(final_dir, exist_ok=True)
        temp_output_pattern = os.path.join(final_dir, f"temp_{start_timestamp}_%02d.ts")
        ffmpeg_command = ["ffmpeg", "-i", source_mp4_path, "-c", "copy", "-map", "0", "-f", "segment", "-segment_time", "10", "-segment_format", "mpegts", "-reset_timestamps", "1", temp_output_pattern]
        subprocess.run(ffmpeg_command, check=True, capture_output=True, text=True)
        
        temp_files = sorted(glob.glob(os.path.join(final_dir, f"temp_{start_timestamp}_*.ts")))
        for i, temp_file_path in enumerate(temp_files):
            current_segment_timestamp = start_timestamp + (i * 10)
            new_filename = datetime.datetime.fromtimestamp(current_segment_timestamp).strftime('%H%M%S') + ".ts"
            new_file_path = os.path.join(final_dir, new_filename)
            os.rename(temp_file_path, new_file_path)
            final_ts_paths.append(new_file_path)
            
        print(f"[流复制成功] {os.path.basename(source_mp4_path)} 已分割为TS")
        return final_ts_paths
    except Exception as e:
        for f in glob.glob(os.path.join(final_dir, f"temp_{start_timestamp}_*.ts")): os.remove(f)
        print(f"[流复制错误] {os.path.basename(source_mp4_path)} 分割异常: {e.stderr if hasattr(e, 'stderr') else e}")
        return []

def create_m3u8_playlist(ts_file_list, camera_code, start_timestamp):
    """为TS文件列表创建M3U8播放列表。"""
    if not ts_file_list: return
    
    print(f"--- [文件操作] {start_timestamp}: 生成M3U8 ---")
    try:
        start_dt = datetime.datetime.fromtimestamp(start_timestamp)
        dest_dir = os.path.join(FINAL_TS_ROOT, camera_code, start_dt.strftime('%Y/%m/%d'), 'source')
        os.makedirs(dest_dir, exist_ok=True)
        
        m3u8_filename = start_dt.strftime('%H%M%S') + "_60.m3u8"
        m3u8_filepath = os.path.join(dest_dir, m3u8_filename)
        
        m3u8_content = ["#EXTM3U", "#EXT-X-VERSION:3", "#EXT-X-TARGETDURATION:10", "#EXT-X-MEDIA-SEQUENCE:0"]
        for ts_path in ts_file_list:
            m3u8_content.append("#EXTINF:10.000,")
            relative_path = f"../{os.path.basename(os.path.dirname(ts_path))}/{os.path.basename(ts_path)}"
            m3u8_content.append(relative_path)
        m3u8_content.append("#EXT-X-ENDLIST")
        
        with open(m3u8_filepath, 'w') as f: f.write('\n'.join(m3u8_content))
        print(f"[文件操作成功] M3U8已保存到: {m3u8_filepath}")
    except Exception as e:
        print(f"[文件操作错误] M3U8生成异常: {e}")

def process_video_task(video_path, camera_code, timestamp):
    """一个完整的视频后处理流程，在一个独立的进程中执行。"""
    print(f"\n>>>>>> 开始在后台进程处理视频: {video_path} <<<<<<")
    
    # MODIFIED: 根据您的要求，已禁用将视频逐帧保存为图片的功能
    # 1. GPU 抽帧 (JPG格式) - 已禁用
    # extract_images_gpu(video_path, camera_code, timestamp)
    
    # 1. GPU 高效生成缩略图 (直接从视频)
    create_jpg_thumbnail_from_video_gpu(video_path, camera_code, timestamp)
    
    # 2. 切片为TS (此过程为流复制，CPU占用低)
    ts_file_list = convert_mp4_to_ts(video_path, camera_code, timestamp)

    # 3. 如果切片成功，则创建M3U8
    if ts_file_list:
        create_m3u8_playlist(ts_file_list, camera_code, timestamp)

    # 4. 清理临时的MP4文件
    try:
        if os.path.exists(video_path):
            os.remove(video_path)
            print(f"\n[任务完成] 已清理临时文件: {video_path}")
    except OSError as e:
        print(f"\n[错误] 清理临时文件失败: {video_path},原因: {e}")
    
    print(f"\n>>>>>> 视频 {video_path} 的所有处理已完成 <<<<<<")

# ============================================================================
# 4. 并发下载与主循环 (已增强启动检查)
# ============================================================================
def download_and_submit_task(dev_instance, camera, start_time, end_time, process_executor):
    camera_code = camera["code"]
    channel_id = camera["channel_id"]
    print(f"[下载线程-{channel_id}] 开始下载摄像头: {camera_code}")
    
    # NVR IP地址作为子目录，避免多个NVR下载时文件冲突
    nvr_ip_str = NVR_IP.decode('utf-8')
    save_dir = os.path.join(LOCAL_VIDEO_SAVE_DIR, nvr_ip_str)
    os.makedirs(save_dir, exist_ok=True)
    
    temp_video_filename = os.path.join(save_dir, f"{camera_code}_{start_time}.mp4")
    
    if dev_instance.download(channel_id, start_time, end_time, temp_video_filename):
        print(f"[下载线程-{channel_id}] 成功: {os.path.basename(temp_video_filename)} 已下载，提交到后台GPU处理。")
        process_executor.submit(process_video_task, temp_video_filename, camera_code, start_time)
    else:
        print(f"[下载线程-{channel_id}] 失败: 摄像头 {camera_code} 下载失败。")

if __name__ == '__main__':
    # sanity check: 检查ffmpeg是否存在
    try:
        result = subprocess.run(["ffmpeg", "-hwaccels"], capture_output=True, text=True, check=True)
        if "nvdec" not in result.stdout:
            print("\n\n!!!! 警告 !!!!")
            print("您的ffmpeg版本可能不支持 'nvdec'。请确保已安装支持NVIDIA硬件加速的ffmpeg。")
            print("程序将继续运行，但可能无法使用GPU加速。")
            print("!!!! 警告 !!!!\n\n")
    except FileNotFoundError:
        print("错误：找不到 'ffmpeg' 命令。请确保ffmpeg已安装并在您的系统路径中。")
        sys.exit(1)

    # 增强了SDK加载、初始化和登录的检查流程
    dev = DevClass()
    if not dev.egsSdk:
        print("因SDK动态库加载失败，程序无法继续运行。")
        sys.exit(1)

    if not dev.sdk_init():
        print("因SDK初始化失败，程序无法继续运行。")
        sys.exit(1)
    
    if not dev.login(NVR_IP, NVR_PORT, NVR_USERNAME, NVR_PASSWORD):
        print("因NVR登录失败，程序无法继续运行。")
        sys.exit(1)

    start_dt = datetime.datetime(2025, 7, 22, 15, 20, 0)
    start_time = int(start_dt.timestamp())

    print(f"\n任务启动，起始时间: {start_dt.strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"模式: 多线程下载 + GPU加速处理 (最多 {MAX_PROCESSING_WORKERS} 个并行进程)")

    with concurrent.futures.ProcessPoolExecutor(max_workers=MAX_PROCESSING_WORKERS) as executor:
        # 设定一个总的下载轮次, 您可以按需修改，例如改为一个很大的数或者while True循环
        for i in range(3):
            try:
                end_time = start_time + DOWNLOAD_DURATION_SECONDS
                print(f"\n--- [主线程-下载轮次 {i+1}/3] 时间片: {datetime.datetime.fromtimestamp(start_time)} ---")
                
                download_threads = []
                for camera in CAMERAS:
                    thread = threading.Thread(target=download_and_submit_task, args=(dev, camera, start_time, end_time, executor))
                    download_threads.append(thread)
                    thread.start()

                for thread in download_threads:
                    thread.join()
                
                # 更新下一次下载的开始时间
                start_time = end_time

            except KeyboardInterrupt:
                print("\n接收到手动停止命令，正在退出...")
                break
            except Exception as e:
                print(f"主循环发生未知错误: {e}。")
                break
            
    print("\n\n所有任务已提交，等待后台处理完成...")
    print("程序已退出。")